"""Split a parallel file into train and dev corpus"""
import io
import random


def split_file(fn, n):
    """Split pairs file into train pairs and dev pairs"""
    print("Reading corpus...")
    pairs = io.open(fn, encoding="utf-8").readlines()

    npairs = len(pairs)
    print(npairs, "pairs")

    print("Shuffling...")
    random.shuffle(pairs)

    print("Splitting...")
    dev_pairs = pairs[:n]
    train_pairs = pairs[n:]

    return train_pairs, dev_pairs


def split_pair(pairs, fn1, fn2):
    """Split pairs file into source file and target file"""
    f1 = io.open(fn1, "w", encoding="utf-8")
    f2 = io.open(fn2, "w", encoding="utf-8")

    cnt = 0
    for p in pairs:
        st = p.strip().split("\t")
        if len(st) >= 2:
            cnt += 1
            if cnt % 5000 == 0:
                print(cnt)

            f1.write(st[0] + "\n")
            f2.write(st[1] + "\n")
    print(cnt)


corpus_file = "big/data.tsv"
dev_n = 1500

train_pairs, dev_pairs = split_file(corpus_file, dev_n)

src_train_file = "big/src-train.txt"
tgt_train_file = "big/tgt-train.txt"

src_dev_file = "big/src-dev.txt"
tgt_dev_file = "big/tgt-dev.txt"

split_pair(train_pairs, src_train_file, tgt_train_file)
split_pair(dev_pairs, src_dev_file, tgt_dev_file)
